# 环节3：文本向量化 - 把文字片段转成"数字密码"
import os

from langchain_community.embeddings import HuggingFaceEmbeddings, DashScopeEmbeddings

# 导入需要的库


from langchain_rag.rag.document_loading import load_document
from langchain_rag.rag.document_splitting import split_documents


# 导入前两个环节的函数



def initialize_embeddings(model_name="all-MiniLM-L6-v2"):
    """
    初始化文本向量化工具

    参数:
        model_name: 预训练模型名称

    返回:
        embeddings: 初始化好的向量化工具
    """
    # 创建HuggingFaceEmbeddings对象，用于文本向量化
    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        # 模型缓存目录，避免重复下载
        cache_folder="./model_cache"
    )
    # embeddings = DashScopeEmbeddings(
    #     model="text-embedding-v4",
    #     dashscope_api_key=os.getenv("DASH_SCOPE_API_KEY")
    # )

    print(f"成功加载向量化模型: {model_name}")
    return embeddings


def vectorize_documents(embeddings, documents):
    """
    将文档片段转换为向量

    参数:
        embeddings: 向量化工具
        documents: 分割后的文档片段列表

    返回:
        vectors: 文档片段对应的向量列表
    """
    # 提取文档内容
    texts = [doc.page_content for doc in documents]

    # 将文本转换为向量
    vectors = embeddings.embed_documents(texts)

    print(f"文本向量化完成！共处理 {len(vectors)} 个片段，每个向量维度为 {len(vectors[0])}")

    return vectors


# 主程序：测试文本向量化功能
if __name__ == "__main__":
    # 1. 加载PDF文档
    pdf_file_path = "平安保险用户手册.pdf"
    document_pages = load_document(pdf_file_path)

    if document_pages:
        # 2. 分割文档
        split_docs = split_documents(document_pages)

        # 3. 初始化向量化工具
        embeddings = initialize_embeddings()

        # 4. 对前5个文档片段进行向量化（全部向量化可能耗时）
        sample_docs = split_docs[:5]
        sample_vectors = vectorize_documents(embeddings, sample_docs)

        # 5. 展示向量化结果
        print("\n----- 向量化结果预览 -----")
        for i, (doc, vec) in enumerate(zip(sample_docs, sample_vectors)):
            print(f"\n片段 {i + 1} 内容预览: {doc.page_content[:50]}...")
            print(f"对应向量（前10个数字）: {vec[:10]}...")  # 只显示前10个维度
            print(f"向量总维度: {len(vec)}")